import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from CSPNet import CSPResNet18
from tqdm import tqdm
import matplotlib.pyplot as plt
import os

# 超参数
batch_size = 128
epochs = 5
lr = 0.1

# 数据预处理
transform = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor()
])

train_dataset = datasets.CIFAR100(root='./data', train=True, transform=transform, download=True)
test_dataset = datasets.CIFAR100(root='./data', train=False, transform=transforms.ToTensor())

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)

# 模型与优化器
device = torch.device('cpu')
model = CSPResNet18().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

# 记录
train_losses, test_losses, accs, top5_accs = [], [], [], []

def evaluate():
    model.eval()
    correct, total, top5_correct = 0, 0, 0
    test_loss = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += criterion(output, target).item()
            _, pred = output.topk(5, 1, True, True)
            total += target.size(0)
            correct += pred[:, :1].eq(target.view(-1, 1)).sum().item()
            top5_correct += pred.eq(target.view(-1, 1).expand_as(pred)).sum().item()
    acc = 100. * correct / total
    top5 = 100. * top5_correct / total
    return test_loss / len(test_loader), acc, top5


for epoch in range(epochs):
    model.train()
    loop = tqdm(train_loader, desc=f'Epoch {epoch+1}/{epochs} [Train]')
    total_loss = 0
    for data, target in loop:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        loop.set_postfix(loss=loss.item())
    scheduler.step()

    # 验证
    val_loss, acc, top5 = evaluate()
    train_losses.append(total_loss / len(train_loader))
    test_losses.append(val_loss)
    accs.append(acc)
    top5_accs.append(top5)

    print(f'Epoch {epoch+1}: val_loss={val_loss:.4f}, acc={acc:.2f}%, top5={top5:.2f}%')


# 保存模型
torch.save(model.state_dict(), 'cspresnet18_cifar100.pth')

# 保存训练曲线
plt.figure()
plt.plot(train_losses, label='Train Loss')
plt.plot(test_losses, label='Test Loss')
plt.legend()
plt.title("Loss Curve")
plt.savefig("loss_curve.png")

plt.figure()
plt.plot(accs, label='Top-1 Accuracy')
plt.plot(top5_accs, label='Top-5 Accuracy')
plt.legend()
plt.title("Accuracy Curve")
plt.savefig("accuracy_curve.png")
